/* Copyright 2018-2023 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=======================================================================*/

// This class has been generated, DO NOT EDIT!

package org.tensorflow.op.tpu;

import java.util.Arrays;
import org.tensorflow.GraphOperation;
import org.tensorflow.Operand;
import org.tensorflow.Operation;
import org.tensorflow.OperationBuilder;
import org.tensorflow.Output;
import org.tensorflow.op.RawOp;
import org.tensorflow.op.RawOpInputs;
import org.tensorflow.op.Scope;
import org.tensorflow.op.annotation.Endpoint;
import org.tensorflow.op.annotation.OpInputsMetadata;
import org.tensorflow.op.annotation.OpMetadata;
import org.tensorflow.op.annotation.Operator;
import org.tensorflow.types.TInt32;

/**
 * An op computes tuple mask of deduplication data from embedding core.
 * The deduplication data receiving from embedding core is a Tensor with
 * type=DT_VARIANT. The tensor itself is an XLA nested tuple, whose elements are
 * rank 1 tensors. This op is to represents types and length of these elements.
 */
@OpMetadata(
    opType = ComputeDedupDataTupleMask.OP_NAME,
    inputsClass = ComputeDedupDataTupleMask.Inputs.class
)
@Operator(
    group = "tpu"
)
public final class ComputeDedupDataTupleMask extends RawOp implements Operand<TInt32> {
  /**
   * The name of this op, as known by TensorFlow core engine
   */
  public static final String OP_NAME = "ComputeDedupDataTupleMask";

  private Output<TInt32> outputShape;

  public ComputeDedupDataTupleMask(Operation operation) {
    super(operation, OP_NAME);
    int outputIdx = 0;
    outputShape = operation.output(outputIdx++);
  }

  /**
   * Factory method to create a class wrapping a new ComputeDedupDataTupleMask operation.
   *
   * @param scope current scope
   * @param config Serialized TPUEmbeddingConfiguration proto.
   * @return a new instance of ComputeDedupDataTupleMask
   */
  @Endpoint(
      describeByClass = true
  )
  public static ComputeDedupDataTupleMask create(Scope scope, String config) {
    OperationBuilder opBuilder = scope.opBuilder(OP_NAME, "ComputeDedupDataTupleMask");
    opBuilder.setAttr("config", config);
    return new ComputeDedupDataTupleMask(opBuilder.build());
  }

  /**
   * Gets outputShape.
   * A 2-D int tensor represent mask of deduplication data tuple generated by
   * {@code XlaRecvTPUEmbeddingDeduplicationData}. The tuple has several integer and float
   * type 1-D tensor tuple elements. The first dimenion of this output_shape 2-D
   * tensor is tensor type of tuple elements, {@code 0} represents integer tensor, {@code 1}
   * represents float tensor. The second dimension of {@code output_shape} gives length of
   * each tuple element.
   * @return outputShape.
   */
  public Output<TInt32> outputShape() {
    return outputShape;
  }

  @Override
  public Output<TInt32> asOutput() {
    return outputShape;
  }

  @OpInputsMetadata(
      outputsClass = ComputeDedupDataTupleMask.class
  )
  public static class Inputs extends RawOpInputs<ComputeDedupDataTupleMask> {
    /**
     * Serialized TPUEmbeddingConfiguration proto.
     */
    public final String config;

    public Inputs(GraphOperation op) {
      super(new ComputeDedupDataTupleMask(op), op, Arrays.asList("config"));
      int inputIndex = 0;
      config = op.attributes().getAttrString("config");
    }
  }
}
